# FROM nvidia/cuda:11.6.1-cudnn8-devel-ubuntu20.04
FROM nvcr.io/nvidia/tensorrt:22.08-py3
ARG CERES_VERSION=2.1.0
ARG CMAKE_VERSION=3.23.1
ARG ONNX_VERSION=1.12.1
ARG OPENCV_VERSION=4.6.0
ENV SWARM_WS=/root/swarm_ws
ARG ROS_VERSION=noetic

#Some basic dependencies
RUN  apt-get -y update && \
      DEBIAN_FRONTEND=noninteractive TZ=Asia/Beijing apt-get -y install tzdata && \
      apt-get install -y wget curl lsb-release git \
      libatlas-base-dev \
      libeigen3-dev \
      libgoogle-glog-dev \
      libsuitesparse-dev \
      libglib2.0-dev \
      libyaml-cpp-dev \
      libdw-dev

#Install ROS
# update ros repository
# some code copied from https://github.com/HKUST-Aerial-Robotics/VINS-Fusion/blob/master/docker/Dockerfile
# RUN sh -c 'echo "deb http://packages.ros.org/ros/ubuntu $(lsb_release -sc) main" > /etc/apt/sources.list.d/ros-latest.list' &&\
RUN sh -c 'echo "deb http://mirrors.ustc.edu.cn/ros/ubuntu/ `lsb_release -cs` main" > /etc/apt/sources.list.d/ros-latest.list' && \
      curl -s https://raw.githubusercontent.com/ros/rosdistro/master/ros.asc | apt-key add - && \
      apt-get update && \
      if [ "x$(nproc)" = "x1" ] ; then export USE_PROC=1 ; \
      else export USE_PROC=$(($(nproc)/2)) ; fi && \
      apt-get update && \
      DEBIAN_FRONTEND=noninteractive  apt-get install -y \
      ros-${ROS_VERSION}-ros-base \
      ros-${ROS_VERSION}-nav-msgs \
      ros-${ROS_VERSION}-sensor-msgs \
      ros-${ROS_VERSION}-cv-bridge \
      ros-${ROS_VERSION}-rviz \
      ros-${ROS_VERSION}-pcl-ros \
      ros-${ROS_VERSION}-image-transport-plugins \
      python3-rosdep \
      python3-rosinstall \
      python3-rosinstall-generator \
      python3-wstool \
      build-essential \
      python3-rosdep \
      ros-${ROS_VERSION}-catkin \
      net-tools \
      python3-catkin-tools \
      htop \
      xterm \
      gdb && \
      rosdep init && \
      rosdep update

#Replace CMAKE
RUN   rm /usr/bin/cmake && \
      wget https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-Linux-x86_64.sh \
      -q -O /tmp/cmake-install.sh \
      && chmod u+x /tmp/cmake-install.sh \
      && /tmp/cmake-install.sh --skip-license --prefix=/usr/ \
      && rm /tmp/cmake-install.sh \
      && cmake --version

#Install ceres
RUN   git clone https://github.com/HKUST-Swarm/ceres-solver -b D2SLAM && \
      cd ceres-solver && \
      mkdir build && cd build && \
      cmake  -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=OFF -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARKS=OFF .. && \
      make -j$(USE_PROC) install && \
      rm -rf ../../ceres-solver && \
      apt-get clean all

RUN   wget https://github.com/microsoft/onnxruntime/releases/download/v${ONNX_VERSION}/onnxruntime-linux-x64-gpu-${ONNX_VERSION}.tgz && \
      tar -zxvf onnxruntime-linux-x64-gpu-${ONNX_VERSION}.tgz && \
      cp -r onnxruntime-linux-x64-gpu-${ONNX_VERSION}/include/* /usr/local/include/ && \
      cp -r onnxruntime-linux-x64-gpu-${ONNX_VERSION}/lib/* /usr/local/lib/ && \
      rm -rf onnxruntime-linux-x64-gpu-${ONNX_VERSION} && \
      rm onnxruntime-linux-x64-gpu-${ONNX_VERSION}.tgz

#Install Libtorch (CUDA)
RUN   wget https://download.pytorch.org/libtorch/nightly/cpu/libtorch-cxx11-abi-shared-with-deps-latest.zip && \
      unzip libtorch-cxx11-abi-shared-with-deps-latest.zip && \
      rm libtorch-cxx11-abi-shared-with-deps-latest.zip && \
      cp -r libtorch/* /usr/local/

#Install LCM
RUN   git clone https://github.com/lcm-proj/lcm && \
      cd lcm && \
      git checkout tags/v1.4.0 && \
      mkdir build && cd build && \
      cmake -DCMAKE_BUILD_TYPE=Release -DBUILD_TESTING=OFF -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARKS=OFF .. && \
      make -j$(USE_PROC) install && \
      pip install lcm

#Install Faiss
RUN   git clone https://github.com/facebookresearch/faiss.git && \
      cd faiss && \
      cmake -B build -DCMAKE_BUILD_TYPE=Release -DFAISS_OPT_LEVEL=avx2 -DFAISS_ENABLE_PYTHON=OFF -DBUILD_TESTING=OFF -DFAISS_ENABLE_GPU=OFF . && \
      make -C build -j faiss && \
      make -C build install && \
      rm -rf ../faiss

#Install OpenCV4 with CUDA
RUN   apt update && \
      apt install libgtk2.0-dev -y && \
      wget https://github.com/opencv/opencv/archive/${OPENCV_VERSION}.zip -O opencv.zip && \
      unzip opencv.zip && \
      rm opencv.zip && \
      git clone https://github.com/opencv/opencv_contrib.git -b ${OPENCV_VERSION}
RUN   cd opencv-${OPENCV_VERSION} && \
      mkdir build && cd build && \
      cmake .. \
            -D CMAKE_BUILD_TYPE=RELEASE \
            -D CMAKE_INSTALL_PREFIX=/usr/local \
            -D WITH_CUDA=ON \
            -D WITH_CUDNN=ON \
            -D WITH_CUBLAS=ON \
            -D WITH_TBB=ON \
            -D OPENCV_DNN_CUDA=ON \
            -D OPENCV_ENABLE_NONFREE=ON \
            -D OPENCV_EXTRA_MODULES_PATH=../../opencv_contrib/modules \
            -D BUILD_EXAMPLES=OFF \
            -D BUILD_opencv_java=OFF \
            -D BUILD_opencv_python=OFF \
            -D BUILD_TESTS=OFF \
            -D BUILD_PERF_TESTS=OFF \
            -D BUILD_opencv_apps=OFF \
            -D BUILD_LIST=calib3d,features2d,highgui,dnn,imgproc,imgcodecs,\
cudev,cudaoptflow,cudaimgproc,cudalegacy,cudaarithm,cudacodec,cudastereo,\
cudafeatures2d,xfeatures2d,tracking,stereo,\
aruco,videoio,ccalib && \
      make -j $(USE_PROC) && \
      make install 

#Install OpenGV
RUN git clone https://github.com/HKUST-Swarm/opengv && \
      mkdir opengv/build && cd opengv/build && cmake .. && make && \
      make install

#Install Backward
RUN   wget https://raw.githubusercontent.com/bombela/backward-cpp/master/backward.hpp -O /usr/local/include/backward.hpp

#Install Vulkan
RUN   wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | sudo apt-key add - && \
      wget -qO /etc/apt/sources.list.d/lunarg-vulkan-1.3.231-focal.list https://packages.lunarg.com/vulkan/1.3.231/lunarg-vulkan-1.3.231-focal.list && \
      apt update && \
      apt install vulkan-sdk -y

#Install taichi
RUN    pip install taichi transformations numpy lcm matplotlib scipy && apt install ros-${ROS_VERSION}-ros-numpy

#Build D2SLAM
RUN   mkdir -p ${SWARM_WS}/src/ && \
      cd ${SWARM_WS}/src/ && \
      git clone https://github.com/HKUST-Swarm/swarm_msgs.git -b D2SLAM && \
      git clone https://github.com/HKUST-Swarm/sync_bag_player.git && \
      git clone https://github.com/ros-perception/vision_opencv.git -b ${ROS_VERSION}
COPY ./ ${SWARM_WS}/src/D2SLAM
WORKDIR $SWARM_WS
RUN   source "/opt/ros/${ROS_VERSION}/setup.bash" && \
      catkin config -DCMAKE_BUILD_TYPE=Release \
            --cmake-args -DONNXRUNTIME_LIB_DIR=/usr/local/lib \
            -DONNXRUNTIME_INC_DIR=/usr/local/include && \
      catkin build
